import pandas as pd
from mlxtend.preprocessing import TransactionEncoder
from mlxtend.frequent_patterns import apriori, association_rules
import matplotlib.pyplot as plt
import networkx as nx

# 设置字体
plt.rcParams['font.sans-serif'] = ['SimHei']  # 中文显示
plt.rcParams['axes.unicode_minus'] = False  # 负号显示

# 读取 Excel 文件中的数据
df = pd.read_excel('C:/关联规则/餐厅数据.xlsx')

# 检查列名
print(df.columns)

# 转换菜品格式，假设‘菜品’列为每个交易的菜品列表
transactions = df['菜品'].str.split(',').to_list()
print(transactions[:5])  # 打印前五个交易数据

# 标准化交易数据
te = TransactionEncoder()
te_ary = te.fit(transactions).transform(transactions)

# 将转换后的数据转换为 DataFrame
encoded_df = pd.DataFrame(te_ary, columns=te.columns_)
print(encoded_df.head())  # 打印转换后的数据

# 使用 apriori 算法进行频繁项集分析
frequent_itemsets = apriori(encoded_df, min_support=0.1, use_colnames=True)

# 排序频繁项集
frequent_itemsets.sort_values(by='support', ascending=False, inplace=True)

# 选择二项集查看
print(frequent_itemsets[frequent_itemsets['itemsets'].apply(lambda x: len(x) == 2)])

# 生成关联规则
rules = association_rules(frequent_itemsets, metric="lift", min_threshold=1)
rules.sort_values(by=['lift', 'confidence'], ascending=False, inplace=True)

# 打印关联规则
print(rules[['antecedents', 'consequents', 'support', 'confidence', 'lift']])

# 可视化关联规则
G = nx.DiGraph()

# 将关联规则添加到图中
for _, row in rules.iterrows():
    antecedent = ', '.join(list(row['antecedents']))  # 转换为字符串
    consequent = ', '.join(list(row['consequents']))  # 转换为字符串
    lift = row['lift']
    confidence = row['confidence']
    
    # 添加边（关联规则：前件->后件），权重为lift
    G.add_edge(antecedent, consequent, weight=lift, confidence=confidence)

# 绘制图形
plt.figure(figsize=(12, 12))

# 计算节点位置
pos = nx.spring_layout(G, k=0.2, iterations=20)

# 绘制网络图
nx.draw_networkx_nodes(G, pos, node_size=500, node_color='lightblue', alpha=0.6)
nx.draw_networkx_edges(G, pos, width=1.0, alpha=0.5, edge_color='gray')
nx.draw_networkx_labels(G, pos, font_size=12, font_weight='bold')

# 边的标签（显示lift值）
edge_labels = nx.get_edge_attributes(G, 'weight')
nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels)

plt.title('关联规则网络图')
plt.axis('off')
plt.show()